## resample all data 

import os 
import numpy as np 
import SimpleITK as sitk 

def resample_img(
    image: sitk.Image,
    out_spacing = (2.0, 2.0, 2.0),
    out_size = None,
    is_label: bool = False,
    pad_value = 0.,
) -> sitk.Image:
    """
    Resample images to target resolution spacing
    Ref: SimpleITK
    """
    # get original spacing and size
    original_spacing = image.GetSpacing()
    original_size = image.GetSize()

    # convert our z, y, x convention to SimpleITK's convention
    out_spacing = list(out_spacing)[::-1]

    if out_size is None:
        # calculate output size in voxels
        out_size = [
            int(np.round(
                size * (spacing_in / spacing_out)
            ))
            for size, spacing_in, spacing_out in zip(original_size, original_spacing, out_spacing)
        ]

    # determine pad value
    if pad_value is None:
        pad_value = image.GetPixelIDValue()

    # set up resampler
    resample = sitk.ResampleImageFilter()
    resample.SetOutputSpacing(list(out_spacing))
    resample.SetSize(out_size)
    resample.SetOutputDirection(image.GetDirection())
    resample.SetOutputOrigin(image.GetOrigin())
    resample.SetTransform(sitk.Transform())
    resample.SetDefaultPixelValue(pad_value)
    if is_label:
        resample.SetInterpolator(sitk.sitkNearestNeighbor)
    else:
        resample.SetInterpolator(sitk.sitkBSpline)

    image = resample.Execute(image)

    return image

def resample_all_data(data_dir, out_dir, out_space, is_label=False):
    all_path = os.listdir(data_dir)
    assert data_dir != out_dir
    os.makedirs(out_dir,exist_ok=True)
    
    for p in all_path:
        input_data = os.path.join(data_dir, p)
        input_data = sitk.ReadImage(input_data)
        
        out_data = resample_img(input_data, out_spacing=out_space, is_label=is_label)
        
        out_arr = sitk.GetArrayFromImage(out_data)
        sitk.WriteImage(out_data, os.path.join(out_dir, p))

        print(f"数据写入成功: {os.path.join(out_dir, p)}, {out_arr.shape}")
        


def resample(data_dir, output_dir):
    train_image_dir = os.path.join(data_dir, "imagesTr")
    train_label_dir = os.path.join(data_dir, "labelsTr")
    val_image_dir = os.path.join(data_dir, "imagesVal")
    val_label_dir = os.path.join(data_dir, "labelsVal")
    test_image_dir = os.path.join(data_dir, "imagesTs")
    test_label_dir = os.path.join(data_dir, "labelsTs")
    
    train_image_dir_save = os.path.join(output_dir, "imagesTr")
    train_label_dir_save = os.path.join(output_dir, "labelsTr")
    val_image_dir_save = os.path.join(output_dir, "imagesVal")
    val_label_dir_save = os.path.join(output_dir, "labelsVal")
    test_image_dir_save = os.path.join(output_dir, "imagesTs")
    test_label_dir_save = os.path.join(output_dir, "labelsTs")
    
    space = [2.5, 2.0, 2.0]
    
    resample_all_data(train_image_dir, train_image_dir_save, out_space=space, is_label=False )
    resample_all_data(train_label_dir, train_label_dir_save, out_space=space, is_label=True )
    
    resample_all_data(test_image_dir, test_image_dir_save, out_space=space, is_label=False )
    resample_all_data(test_label_dir, test_label_dir_save, out_space=space, is_label=True )
    
    
if __name__ == "__main__":
    resample(data_dir="/home/xingzhaohu/sharefs/datasets/WORD-V0.1.0", output_dir="/home/xingzhaohu/sharefs/datasets/WORD-V0.1.0-resample2.0")